package org.mx;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;

import org.mx.sm.SM4;
import org.mx.sm.SM4Context;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 数据加解密工具类
 *
 * @author : john.peng date : 2020/3/26
 */
public class SecretUtils {
	private static Logger logger = LoggerFactory.getLogger(SecretUtils.class);

	/**
	 * 默认的构造函数
	 */
	private SecretUtils() {
		super();
	}

	/**
	 * 采用56位Key的DES算法加密数据
	 * 
	 * @param password 密码
	 * @param src 待加密的数据
	 * 
	 * @return 加密后的数据
	 */
	public static byte[] des56Encrypt(String password, byte[] src) {
		return secretCrypt("DES", 56, true, password, src);
	}

	/**
	 * 采用56位Key的DES算法解密数据
	 * 
	 * @param password 密码
	 * @param src 待解密的数据
	 * 
	 * @return 解密后的数据
	 */
	public static byte[] des56Decrypt(String password, byte[] src) {
		return secretCrypt("DES", 56, false, password, src);
	}

	/**
	 * 采用112位Key的3DES算法加密数据
	 * 
	 * @param password 密码
	 * @param src 待加密的数据
	 * 
	 * @return 加密后的数据
	 */
	public static byte[] desede112Encrypt(String password, byte[] src) {
		return secretCrypt("DESede", 112, true, password, src);
	}

	/**
	 * 采用112位Key的3DES算法解密数据
	 * 
	 * @param password 密码
	 * @param src 待解密的数据
	 * 
	 * @return 解密后的数据
	 */
	public static byte[] desede112Decrypt(String password, byte[] src) {
		return secretCrypt("DESede", 112, false, password, src);
	}

	/**
	 * 采用168位Key的3DES算法加密数据
	 * 
	 * @param password 密码
	 * @param src 待加密的数据
	 * 
	 * @return 加密后的数据
	 */
	public static byte[] desede168Encrypt(String password, byte[] src) {
		return secretCrypt("DESede", 168, true, password, src);
	}

	/**
	 * 采用168位Key的3DES算法解密数据
	 * 
	 * @param password 密码
	 * @param src 待解密的数据
	 * 
	 * @return 解密后的数据
	 */
	public static byte[] desede168Decrypt(String password, byte[] src) {
		return secretCrypt("DESede", 168, false, password, src);
	}

	/**
	 * 采用192位Key的AES算法加密数据
	 * 
	 * @param password 密码
	 * @param src 待加密的数据
	 * 
	 * @return 加密后的数据
	 */
	public static byte[] aes192Encrypt(String password, byte[] src) {
		return secretCrypt("AES", 192, true, password, src);
	}

	/**
	 * 采用192位Key的AES算法解密数据
	 * 
	 * @param password 密码
	 * @param src 待解密的数据
	 * 
	 * @return 解密后的数据
	 */
	public static byte[] aes192Decrypt(String password, byte[] src) {
		return secretCrypt("AES", 192, false, password, src);
	}

	/**
	 * 采用256位Key的AES算法加密数据
	 * 
	 * @param password 密码
	 * @param src 待加密的数据
	 * 
	 * @return 加密后的数据
	 */
	public static byte[] aes256Encrypt(String password, byte[] src) {
		return secretCrypt("AES", 256, true, password, src);
	}

	/**
	 * 采用256位Key的AES算法解密数据
	 * 
	 * @param password 密码
	 * @param src 待解密的数据
	 * 
	 * @return 解密后的数据
	 */
	public static byte[] aes256Decrypt(String password, byte[] src) {
		return secretCrypt("AES", 256, false, password, src);
	}

	/**
	 * 创建指定的SecretKey对象
	 * 
	 * @param algorithm 算法名称，支持：DES、DESede、AES
	 * @param keySize key长度，DES（56位）、DESede（112、168位）、AES（192、256位）
	 * @param password 密码
	 * 
	 * @return SecretKey
	 * 
	 * @throws NoSuchAlgorithmException 不支持的算法
	 */
	private static SecretKey generateSecretKey(String algorithm, int keySize, String password)
			throws NoSuchAlgorithmException {
		KeyGenerator keyGen = KeyGenerator.getInstance(algorithm);
		SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
		random.setSeed(password.getBytes());
		keyGen.init(keySize, random);
		return keyGen.generateKey();
	}

	/**
	 * 支持DES、3DES（DESede）、AES等算法的对称加解密
	 * 
	 * @param algorithm 算法名称，支持：DES、DESede、AES
	 * @param keySize key长度，DES（56位）、DESede（112、168位）、AES（192、256位）
	 * @param encrypted 如果设置为true表示加密，否则为解密
	 * @param password 密码
	 * @param src 处理的数据字节数组
	 * 
	 * @return 处理完成后数据字节数组，如果发生错误，则返回null
	 */
	private static byte[] secretCrypt(String algorithm, int keySize, boolean encrypted, String password, byte[] src) {
		SecretKey skey;
		try {
			skey = generateSecretKey(algorithm, keySize, password);
		} catch (NoSuchAlgorithmException ex) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("Unsupported algorithm: %s.", algorithm), ex);
			}
			return null;
		}
		Cipher cipher;
		try {
			cipher = Cipher.getInstance(algorithm);
		} catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("Unsupported algorithm: %s.", algorithm), ex);
			}
			return null;
		}
		try {
			cipher.init(encrypted ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE, skey);
		} catch (InvalidKeyException ex) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("Invalide key: %s.", password), ex);
			}
			return null;
		}
		try {
			return cipher.doFinal(src);
		} catch (IllegalBlockSizeException | BadPaddingException ex) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("%s the data fail, key: %s.", encrypted ? "Encrypt" : "Decrypt", password),
						ex);
			}
			return null;
		}
	}

	/**
	 * 采用ECB方式的SM4加密方法
	 * 
	 * @param secretKey 密钥
	 * @param src       待加密数据
	 * @return 加密后的数据
	 */
	public static byte[] sm4EncryptECB(String secretKey, byte[] src) {
		try {
			SM4Context ctx = new SM4Context();
			ctx.isPadding = true;
			ctx.mode = SM4.SM4_ENCRYPT;

			byte[] keyBytes = secretKey.getBytes();
			SM4 sm4 = new SM4();
			sm4.sm4_setkey_enc(ctx, keyBytes);
			return sm4.sm4_crypt_ecb(ctx, src);
		} catch (Throwable ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Encrypt any data by SM4(ECB) fail.", ex);
			}
			return null;
		}
	}

	/**
	 * 采用ECB方式的SM4解密方法
	 * 
	 * @param secretKey 密钥
	 * @param src       待解密数据
	 * @return 解密后的数据
	 */
	public static byte[] sm4DecryptECB(String secretKey, byte[] src) {
		try {
			SM4Context ctx = new SM4Context();
			ctx.isPadding = true;
			ctx.mode = SM4.SM4_DECRYPT;
			byte[] keyBytes = secretKey.getBytes();

			SM4 sm4 = new SM4();
			sm4.sm4_setkey_dec(ctx, keyBytes);
			return sm4.sm4_crypt_ecb(ctx, src);
		} catch (Throwable ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Decrypt any data by SM4(ECB) fail.", ex);
			}
			return null;
		}
	}

	/**
	 * 采用CBC方式的SM4加密方法
	 * 
	 * @param secretKey 密钥
	 * @param ivKey     iv密钥
	 * @param src       带加密数据
	 * @return 加密后的数据
	 */
	public static byte[] sm4EncryptCBC(String secretKey, String ivKey, byte[] src) {
		try {
			SM4Context ctx = new SM4Context();
			ctx.isPadding = true;
			ctx.mode = SM4.SM4_ENCRYPT;

			byte[] keyBytes = secretKey.getBytes();
			byte[] ivBytes = ivKey.getBytes();
			SM4 sm4 = new SM4();
			sm4.sm4_setkey_enc(ctx, keyBytes);
			return sm4.sm4_crypt_cbc(ctx, ivBytes, src);
		} catch (Throwable ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Encrypt any data by SM4(CBC) fail.", ex);
			}
			return null;
		}
	}

	/**
	 * 采用CBC方式的SM4解密方法
	 * 
	 * @param secretKey 密钥
	 * @param ivKey     iv密钥
	 * @param src       待解密数据
	 * @return 解密后的数据
	 */
	public static byte[] sm4DecryptCBC(String secretKey, String ivKey, byte[] src) {
		try {
			SM4Context ctx = new SM4Context();
			ctx.isPadding = true;
			ctx.mode = SM4.SM4_DECRYPT;

			byte[] keyBytes = secretKey.getBytes();
			byte[] ivBytes = ivKey.getBytes();
			SM4 sm4 = new SM4();
			sm4.sm4_setkey_dec(ctx, keyBytes);
			return sm4.sm4_crypt_cbc(ctx, ivBytes, src);
		} catch (Throwable ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Decrypt any data by SM4(CBC) fail.", ex);
			}
			return null;
		}
	}
}
